package com.cryptoregistry.compat;

import java.io.File;
import java.io.FileInputStream;
import java.io.FileOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.security.KeyStore;
import java.security.KeyStoreException;
import java.security.cert.CertificateFactory;
import java.security.cert.PKIXParameters;
import java.security.cert.TrustAnchor;
import java.security.cert.X509Certificate;
import java.util.ArrayList;
import java.util.Collection;
import java.util.Iterator;
import java.util.List;

import org.apache.commons.io.FileUtils;
import org.apache.commons.io.FilenameUtils;
import org.apache.commons.io.filefilter.NameFileFilter;
import org.apache.commons.io.filefilter.TrueFileFilter;

import com.cryptoregistry.compat.CmdLineParser.Option;
import com.cryptoregistry.compat.CmdLineParser.OptionException;

import javax.naming.ldap.LdapName;


/**
 * Given a root path, scan for cacert files and install our custom ones.
 * 
 * For example: scanner "C:\Program Files\Java" C:\mycerts
 * 
 * @author Dave
 *
 */
public class CacertScanner {

	String rootPath; // scanner looks here
	String certPath; // location of x509 certs to install
	
	List<CertContainer> customCerts;
	List<X509Certificate> currentStoreCerts;
	
	char [] password = "changeit".toCharArray();
	String keystoreFileName = "cacerts";
	
	boolean verbose, noaction;

	/**
	 * Use with cacerts and assume the default password "changeit"
	 * 
	 * @param rootPath
	 * @param certPath
	 */
	public CacertScanner(String rootPath, String certPath) {
		this.rootPath = rootPath;
		this.certPath = certPath;
		loadCustomCerts();
		currentStoreCerts = new ArrayList<X509Certificate>();
	}
	
	/**
	 * Use when updating a keystore (or keystores of the same name in different folders) other than cacerts. The password
	 * needs to be set for that to work (unless it is also 'changeit' which would be bad
	 * 
	 * @param rootPath
	 * @param certPath
	 * @param password
	 * @param keystoreName
	 */
	public CacertScanner(String rootPath, String certPath, char[] password, String keystoreName) {
		this(rootPath,certPath);
		this.password = password;
		this.keystoreFileName = keystoreName;
	}

	public void scan() throws Exception {
		
		File root = new File(rootPath);
		NameFileFilter nameFilter = new NameFileFilter("cacerts");
		Iterator<File> iter = FileUtils.iterateFiles(root, nameFilter, TrueFileFilter.INSTANCE);
		while(iter.hasNext()){
			File file = iter.next();
			System.out.println("Found "+file.getCanonicalPath());
			  FileInputStream is = new FileInputStream(file);
	            KeyStore keystore = KeyStore.getInstance(KeyStore.getDefaultType());
	            keystore.load(is, password);

	            // This class retrieves the CAs from the keystore
	            PKIXParameters params = new PKIXParameters(keystore);

	            // Get the set of trust anchors, which contain the most-trusted CA certificates
	            Iterator<TrustAnchor> it = params.getTrustAnchors().iterator();
	            while( it.hasNext() ) {
	                TrustAnchor ta = (TrustAnchor)it.next();
	                // Get certificate
	                X509Certificate cert = ta.getTrustedCert();
	            if(verbose)    System.out.println("Found "+cert.getSubjectDN().getName()+", serial:"+cert.getSerialNumber());
	               this.currentStoreCerts.add(cert);
	            }
	            
	            // we've collected all the trust store certs for one cacerts file into a list.
	            // check for our custom certs in there. If we don't find the cert, put it on the add list
	            List<CertContainer> addList = new ArrayList<CertContainer>();
	            for(CertContainer certC: customCerts){
	            	boolean found = false;
	            	for(X509Certificate certT: this.currentStoreCerts){
	            		if(equal(certC.cert,certT)){
	            			found= true;
	            			break;
	            		}
	            	}
	            	if(!found) {
	            		addList.add(certC);
	            		System.err.println("Did not find "+certC.fileName+", will add it");
	            	}else{
	            		System.err.println("Found "+certC.fileName+", not adding this one");
	            	}
	            }
	            
	            // addList now contains any certs not found in this cacerts file that we want to add
	            // open cacerts file, add from the add list.
	            
	            if(!noaction){
		            File destBackup = new File(file.getParentFile(), "cacerts.old");
		            boolean success = backupCacerts(file, destBackup);
		            if(!success) {
		            	System.err.println("Fail, we could not backup original cacerts");
		            	return;
		            }else{
		            	addToCacerts(file, addList);
		            }
	            }else{
	            	System.out.println("the noaction flag is set, so no changes made");
	            }
	            
		}
	}
	
	private boolean backupCacerts(File srcCacerts, File destCacerts){
		try {
			FileUtils.copyFile(srcCacerts, destCacerts);
		} catch (IOException e) {
			return false;
		}
		
		return true;
	}
	
	private void addToCacerts(File cacerts, List<CertContainer> list){
		
		if(list.size() == 0) return;
		
		FileInputStream in = null;
		KeyStore keystore = null;
		
		try {
			in = new FileInputStream(cacerts);
			keystore = KeyStore.getInstance(KeyStore.getDefaultType());
			keystore.load(in, password);
		}catch(Exception x){
			x.printStackTrace();
		}finally{
			if(in != null)
				try {
					in.close();
				} catch (IOException e) {}
		}

		for(CertContainer cert: list){
			try {
				keystore.setCertificateEntry(cert.fileName, cert.cert);
				System.out.println("Added: "+cert.fileName+", "+cert.cert.getSubjectDN());
			} catch (KeyStoreException e) {
				e.printStackTrace();
			}
		}
	  
	    FileOutputStream out = null;
	    
	    try {
	    	out = new FileOutputStream(cacerts);
			keystore.store(out, password);
			System.out.println("Saved the updated keystore file: "+cacerts);
	    }catch(Exception x){
	    	x.printStackTrace();
	    }finally{
	    	if(out != null)
				try {
					out.close();
				} catch (IOException e) {}
	    }
	    
	   
	}

	private void loadCustomCerts() {
		
		customCerts = new ArrayList<CertContainer>();
		
		Collection<File> files = FileUtils.listFiles(new File(certPath), null, true);
		System.out.println("Loading our custom certs...found "+files.size()+":");
		for(File item: files){
			InputStream inStream = null;
			try {
				inStream = new FileInputStream(item);
				  CertificateFactory fact = CertificateFactory.getInstance("X.509");
				  X509Certificate cert = (X509Certificate) fact.generateCertificate(inStream);
				 System.out.println("Using "+cert.getSubjectDN().getName()+", serial:"+cert.getSerialNumber());
				 this.customCerts.add(new CertContainer(createAlias(item),cert));
			}catch(Exception x){
				if(inStream != null) {
					 try {
						inStream.close();
					} catch (IOException e) {
						// TODO Auto-generated catch block
						e.printStackTrace();
					}
				}
			} 
		}
	}
	
	/**
	 * Compare DN and serial number also
	 * 
	 * @param cert0
	 * @param cert1
	 * @return
	 * @throws Exception
	 */
	private boolean equal(X509Certificate cert0, X509Certificate cert1) throws Exception{
		String dn0 = cert0.getSubjectX500Principal().getName();
		String dn1 = cert1.getSubjectX500Principal().getName();
		LdapName ldapDN0 = new LdapName(dn0);
		LdapName ldapDN1 = new LdapName(dn1);
		if(!ldapDN0.equals(ldapDN1)) return false;
		if(!cert0.getSerialNumber().toString().equals(cert1.getSerialNumber().toString())) return false;
		
		return true;
	}
	
	private String createAlias(File file){
		return FilenameUtils.getBaseName(file.getName());
	}
	
	class CertContainer {
		final String fileName; // used as the alias when added to keystore
		final X509Certificate cert; // the cert we will add
		
		public CertContainer(String fileName, X509Certificate cert) {
			super();
			this.fileName = fileName;
			this.cert = cert;
		}
		
	}
	
	public void setRootPath(String rootPath) {
		this.rootPath = rootPath;
	}

	public void setCertPath(String certPath) {
		this.certPath = certPath;
	}

	public void setPassword(char[] password) {
		this.password = password;
	}

	public void setKeystoreFileName(String keystoreFileName) {
		this.keystoreFileName = keystoreFileName;
	}

	public void setVerbose(boolean verbose) {
		this.verbose = verbose;
	}

	public void setNoaction(boolean noaction) {
		this.noaction = noaction;
	}

	public static void main(String [] args){
		
		CmdLineParser parser = new CmdLineParser();
		Option<String> searchOpt = parser.addStringOption('s', "searchpath");
		Option<String> customCertsOpt = parser.addStringOption('c', "customcerts");
		Option<Boolean> verboseOpt = parser.addBooleanOption('v', "verbose");
		Option<Boolean> scanOnlyOpt = parser.addBooleanOption('n', "noaction");
		Option<Boolean> helpOpt = parser.addBooleanOption('h', "help");
		

		try {
			parser.parse(args);
		} catch (OptionException e) {
			e.printStackTrace();
		}
		
		boolean help = parser.getOptionValue(helpOpt, false);
		
		if(help) {
			showHelp("/help.txt");
			System.exit(0);
		}
		
		boolean noaction = parser.getOptionValue(scanOnlyOpt, false);
		boolean verbose = parser.getOptionValue(verboseOpt, false);
		
		String searchPath = parser.getOptionValue(searchOpt, "");
		if(searchPath.equals("")) {
			System.err.println("Please set s|searchpath to a directory path");
			System.exit(2);
		}
		
		String custPath = parser.getOptionValue(customCertsOpt, "");
		if(custPath.equals("")) {
			System.err.println("Please set c|customcerts to a directory path where your certs are located");
			System.exit(2);
		}
		
		try {
			CacertScanner scanner = new CacertScanner(searchPath,custPath);
			scanner.setNoaction(noaction);
			scanner.setVerbose(verbose);
			scanner.scan();
		} catch (Exception e) {
			// TODO Auto-generated catch block
			e.printStackTrace();
			System.exit(2);
		}
	}
	
	public static final void showHelp(String path) {
		InputStream in = Thread.currentThread().getClass().getResourceAsStream(path);
		try {
	        byte[] buf = new byte[1024];
	        int len;
	        while((len=in.read(buf))>0){
	            System.out.write(buf,0,len);
	        }
	       
	    } catch (Exception e) {
	        e.printStackTrace();
	    }finally{
	    	if(in != null){
	    		 try {
					in.close();
				} catch (IOException e) {}
	    	}
	    }
	}
	
}
